import { Expose, Type } from 'class-transformer';
import { Role } from './message';
import { Assert } from '@celljs/core';
import { ContentBlock } from './content-block';
import { Usage } from './usage';

/**
 * Chat completion response object.
 */
export class ChatResponse {

    /**
     * Unique object identifier. The format and length of IDs may change over
     * time.
     */
    @Expose()
    id: string;

    /**
     * Object type. For Messages, this is always "message".
     */
    @Expose()
    type: string;

    /**
     * Conversational role of the generated message. This will always be
     * "assistant".
     */
    @Expose()
    role: Role;

    /**
     * Content generated by the model. This is an array of content blocks.
     */
    @Type(() => ContentBlock)
    content: ContentBlock[];

    /**
     * The model that handled the request.
     */
    @Expose()
    model: string;

    /**
     * The reason the model stopped generating tokens. This will be one
     * of "end_turn", "max_tokens", "stop_sequence", "tool_use", or "timeout".
     */
    @Expose({ name: 'stop_reason' })
    stopReason: string;

    /**
     * Which custom stop sequence was generated, if any.
     */
    @Expose({ name: 'stop_sequence' })
    stopSequence?: string;

    /**
     * Input and output token usage.
     */
    @Expose()
    usage: Usage;

    constructor(
        id: string,
        type: string,
        role: Role,
        content: ContentBlock[],
        model: string,
        stopReason: string,
        stopSequence: string | undefined,
        usage: any
    ) {
        this.id = id;
        this.type = type;
        this.role = role;
        this.content = content;
        this.model = model;
        this.stopReason = stopReason;
        this.stopSequence = stopSequence;
        this.usage = usage;
    }

    static builder(): ChatResponseBuilder {
        return new ChatResponseBuilder();
    }
}

export class ChatResponseBuilder {
    private type: string;
    private id: string;
    private role: Role;
    private content: ContentBlock[];
    private model: string;
    private stopReason: string;
    private stopSequence?: string;
    private usage: Usage;

    constructor() {
    }

    withType(type: string): ChatResponseBuilder {
        this.type = type;
        return this;
    }

    withId(id: string): ChatResponseBuilder {
        this.id = id;
        return this;
    }

    withRole(role: Role): ChatResponseBuilder {
        this.role = role;
        return this;
    }

    withContent(content: ContentBlock[]): ChatResponseBuilder {
        this.content = content;
        return this;
    }

    withModel(model: string): ChatResponseBuilder {
        this.model = model;
        return this;
    }

    withStopReason(stopReason: string): ChatResponseBuilder {
        this.stopReason = stopReason;
        return this;
    }

    withStopSequence(stopSequence: string): ChatResponseBuilder {
        this.stopSequence = stopSequence;
        return this;
    }

    withUsage(usage: Usage): ChatResponseBuilder {
        this.usage = usage;
        return this;
    }

    withOutputTokens(outputTokens: number): ChatResponseBuilder {
        Assert.isTrue(!!this.usage, 'Usage must be set before setting output tokens');
        this.usage.outputTokens = outputTokens;
        return this;
    }

    build(): ChatResponse {
        return new ChatResponse(this.id, this.type, this.role, this.content, this.model, this.stopReason, this.stopSequence, this.usage);
    }
}

export class ChatResponseBuilderReference {

    private static instance: ChatResponseBuilder;

    static get(): ChatResponseBuilder {
        if (!this.instance) {
            this.instance = new ChatResponseBuilder();
        }
        return this.instance;
    }

    static reset(): void {
        this.instance = new ChatResponseBuilder();
    }
}
